# -*- coding: utf-8 -*-

import manage
import single
import elp
import casadi
import numpy as np
import muti
import const
import rect

pos_debug = None
R_DEFAULT = 5.0


def GetEscaperWillGo(state, R=2.0):
    opt = casadi.Opti()
    e = elp.Ellipse(opt)
    # opt.minimize(e.a ** (-2) + e.b ** (-2))
    # opt.minimize(- e.a * e.b)
    opt.minimize((e.a * e.b)**(-1))
    # opt.minimize(-(e.x0 - state.e_pos[0])**2 + (e.y0 - state.e_pos[1])**2)
    opt.subject_to(
        (e.x0 - state.e_pos[0])**2 + (e.y0 - state.e_pos[1])**2 <= R**2
    )
    seg = 36
    # sample = [None] * seg
    # for i in range(seg):
    #     sample[i] = e.Sample(np.pi * 2 * i / seg)
    # for i in range(state.p_n):
    #     o, r = muti.ACircle(
    #         state.p_pos[i], state.e_pos, 
    #         state.p_vmax[i] / state.e_vmax
    #     )
    #     for j in range(seg):
    #         x, y = sample[j]
    #         opt.subject_to(
    #             (x - o[0])**2 + (y - o[1])**2 >= r**2
    #         )
    d = 0.01
    for i in range(state.p_n):
        o, r = muti.ACircle(
            state.p_pos[i], state.e_pos, 
            state.p_vmax[i] / state.e_vmax
        )
        a = d/r
        a = np.pi * 2 / a
        a = int(a) + 1
        angle = np.linspace(0, np.pi*2, a)
        pos = np.zeros((a, 2))
        pos[:, 0] = o[0] + np.cos(angle)*r
        pos[:, 1] = o[1] + np.sin(angle)*r
        for j in range(a):
            if np.linalg.norm(pos[j] - state.e_pos) <= R:
                opt.subject_to(e(pos[j]) >= 0)
    o_pos = single.obstacle_points.query_ball_point(
        state.e_pos, R+0.1)
    o_pos = single.obstacle_points.data[o_pos]
    for i in range(o_pos.shape[0]):
        opt.subject_to(e(o_pos[i]) >= 0)
    opt.subject_to(e.a <= 2*R)
    opt.subject_to(e.b > 0)
    e.BeforeSolve()
    opt.set_initial(e.x0, state.e_pos[0])
    opt.set_initial(e.y0, state.e_pos[1])
    opt.solver('ipopt', const.OPTION)
    try:
        sol = opt.solve()
    except BaseException as e:
        print(e)
        return None
    e.AfterSolve(sol)
    e = np.array([e.x0, e.y0])
    return e


def GetOlstRlst(state):
    olst = np.zeros((state.p_n, 2))
    rlst = np.zeros((state.p_n,))
    for i in range(state.p_n):
        o, r = muti.ACircle(
            state.p_pos[i], state.e_pos, 
            state.p_vmax[i] / state.e_vmax
        )
        olst[i] = o
        rlst[i] = r
    return olst, rlst


def GetbOutDistance(state, pos, olst, rlst):
    path = single.GetPath(state.e_pos, pos)
    bOut = True
    for j in range(path.shape[0]):
        for k in range(state.p_n):
            if np.linalg.norm(path[j] - olst[k]) < rlst[k]:
                bOut = False
                break
    if not bOut:
        return False, np.inf
    d = [0.0] * state.p_n
    for k in range(state.p_n):
        d[k] = single.GetPathDistance(state.p_pos[k], pos) / state.p_vmax[k]
    d = min(d)
    return True, d

g_LastEscaperWillGo = None

def GetEscaperWillGo2(state, R=R_DEFAULT):
    pos = single.kdcenter.query_ball_point(
        state.e_pos, R+0.1)
    pos = single.kdcenter.data[pos]
    olst, rlst = GetOlstRlst(state)
    dislst = []
    elst = []
    for i in range(pos.shape[0]):
        bOut, d = GetbOutDistance(state, pos[i], olst, rlst)
        if not bOut:
            continue
        dislst.append(d)
        elst.append(pos[i])
    dislst = np.array(dislst)
    elst = np.array(elst)
    global g_LastEscaperWillGo
    if g_LastEscaperWillGo is None:
        g_LastEscaperWillGo = elst[dislst.argmax()]
        return g_LastEscaperWillGo
    dislst -= dislst.min()
    dislst /= dislst.max() + 0.01
    elst = elst[np.where(dislst >= 0.8)]
    md = np.inf
    ret = None
    for e in elst:
        d = np.linalg.norm(e - g_LastEscaperWillGo)
        if md < d:
            md = d
            ret = e
    g_LastEscaperWillGo = ret
    return ret


def GetPathPoint(cache, path, scale):
    dispath = single.GetPathDistance(path[0], path[-1])
    aim = path[-1]
    dis = scale
    for j in range(path.shape[0]):
        d = single.GetPathDistance(cache.e_pos, path[j])
        if d >= dispath * dis:
            aim = path[j]
            p1 = j-1
            if p1 < 0:
                p1 = 0
            p1 = path[p1]
            p2 = path[j]
            aim = p2 + (p1 - p2) * ((d - dispath * dis) * np.linalg.norm(p1-p2))
            break
    return aim


# GETTARGET_FUNC = single.GetTarget2
GETTARGET_FUNC = rect.GetTarget

class Manager(manage.Manager):

    def MainLoop(self):
        cache = self.InitState()
        e = GetEscaperWillGo2(cache)
        if e is None:
            return
        order = self.GetOrder(cache)
        path = single.GetPath(cache.e_pos, e)
        aim = np.zeros((cache.p_n, 2))
        for i in range(cache.p_n):
            aim[i] = GetPathPoint(cache, path, order[i] / (cache.p_n - 1))
            self.pursers[i].aim = aim.copy()
        self.AvoidCollision(cache, order, aim)
    
    def GetOrder(self, cache):
        dis = np.zeros((cache.p_n, ))
        for i in range(cache.p_n):
            dis[i] = single.GetPathDistance(
                cache.p_pos[i],
                cache.e_pos
            )
        def key(i):
            return dis[i]
        order = list(range(cache.p_n))
        order.sort(key=key)
        return order


    def AvoidCollision(self, cache, order, aim):
        global GETTARGET_FUNC
        subject_points = np.zeros((0, 2))
        in_points = [None] * cache.p_n
        for i in order:
            p = self.pursers[i]
            in_points[i] = p.GetInPoints()
        for i in order:
            p = self.pursers[i]
            self.pursers[i].aim = aim[i].copy()
            sp = subject_points
            for j in order:
                if j == i:
                    continue
                sp = np.vstack((sp, in_points[j]))
            elp, sample = GETTARGET_FUNC(
                cache.p_pos[i],
                aim[i],
                sp,
                p.GetRMax(),
                p.GetRMin(),
                p.GetInPoints()
            )
            if elp is None:
                continue
            self.elpcache[i] = elp
            x, y = sample
            seg = x.shape[0]
            xy = np.zeros((seg, 2))
            xy[:, 0] = x
            xy[:, 1] = y
            subject_points = np.vstack((subject_points, xy))


def GetEscaperWillGo3(state, R=R_DEFAULT):
    pos = single.kdcenter.query_ball_point(
        state.e_pos, R+0.1)
    pos = single.kdcenter.data[pos]
    olst, rlst = GetOlstRlst(state)
    dislst = []
    elst = []
    for i in range(pos.shape[0]):
        bOut, d = GetbOutDistance(state, pos[i], olst, rlst)
        if not bOut:
            continue
        dislst.append(d)
        elst.append(pos[i])
    dislst = np.array(dislst)
    elst = np.array(elst)
    if dislst.shape[0] <= 0:
        return np.zeros((state.p_n, 0, 2))
    dislst -= dislst.min()
    dislst /= dislst.max() + 0.01
    elst = elst[np.where(dislst >= 0.8)]
    table = np.zeros((state.p_n, elst.shape[0], 2))
    for i in range(state.p_n):
        for j in range(elst.shape[0]):
            path = single.GetPath(state.e_pos, elst[j])
            table[i, j] = GetPathPoint(state, path, i / (state.p_n-1))
    return table


class Manager2(Manager):

    def MainLoop(self):
        cache = self.InitState()
        table = GetEscaperWillGo3(cache)
        order = self.GetOrder(cache)
        aims = np.zeros((cache.p_n, 2))
        aims[:] = cache.e_pos
        for i, p in enumerate(order):
            dmin = np.inf
            for j in range(table.shape[1]):
                d = single.GetPathDistance(cache.p_pos[p], table[i, j])
                if d < dmin:
                    dmin = d
                    aims[p] = table[i, j]
            self.pursers[p].aim = aims[p].copy()
        self.AvoidCollision(cache, order, aims)